from asyncio import FastChildWatcher
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2

from draw import *

from pcgworker.PCGWorker import *

from wfc_env import *

'''
simulatiom parameters
'''

num_matrix_envs = 1
headless = False
spd = 5
max_actors_num_per_env = 9*9*4+1

# get default set of parameters
sim_params = gymapi.SimParams()
args = gymutil.parse_arguments()
sim_params.physx.solver_type = 1
sim_params.physx.num_position_iterations = 2
sim_params.physx.num_velocity_iterations = 1
sim_params.physx.num_threads = args.num_threads
sim_params.physx.use_gpu = args.use_gpu
sim_params.use_gpu_pipeline = False #args.use_gpu_pipeline
sim_params.substeps = 2
sim_params.dt = 1.0 / 30.0
# setting up the Z-up axis
sim_params.up_axis = gymapi.UP_AXIS_Z
sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)

# configure the ground plane
plane_params = gymapi.PlaneParams()
plane_params.normal = gymapi.Vec3(0, 0, 1) # z-up!
plane_params.distance = 1
plane_params.static_friction = 1
plane_params.dynamic_friction = 1
plane_params.restitution = 0.1

# set up the env grid
num_envs = 1
spacing = 20
env_lower = gymapi.Vec3(-spacing, -spacing, -spacing)
env_upper = gymapi.Vec3(spacing, spacing, spacing)

# Attractor setup
attractor_properties = gymapi.AttractorProperties()
attractor_properties.stiffness = 5e10
attractor_properties.damping = 0
attractor_properties.axes = gymapi.AXIS_ROTATION
attractor_pose = gymapi.Transform()
attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
attractor_pose.r = gymapi.Quat(0,1,0,1)

# create procedural asset
asset_options = gymapi.AssetOptions()
asset_options.density = 1.0
asset_options.fix_base_link = False
asset_options1 = gymapi.AssetOptions()
asset_options1.density = 0.02
asset_options1.fix_base_link = True
asset_options1.linear_damping = 0.

color_red = gymapi.Vec3(1,0,0)

# load PCG asset
asset_root = "./assets"
pt_gray_cube = "PCG/gray_cube.urdf"
pt_blue_cube = "PCG/blue_cube.urdf"
pt_red_cube = "PCG/red_cube.urdf"
pt_white_cube = "PCG/white_cube.urdf"
pt_yellow_cube = "PCG/yellow_cube.urdf"
pt_blue_ramp = "PCG/blue_ramp.urdf"
pt_red_ramp = "PCG/red_ramp.urdf"
pt_white_ramp = "PCG/white_ramp.urdf"
pt_yellow_ramp = "PCG/yellow_ramp.urdf"

# ramp and corner attitutde
ramp_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0*math.pi)
ramp_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi)
ramp_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1*math.pi)
ramp_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1.5*math.pi)
corner_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0*math.pi)
corner_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi)
corner_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1*math.pi)
corner_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1.5*math.pi)

# birick size
side_l = 2

# camera sensor properties
camera_properties = gymapi.CameraProperties()
camera_properties.width = WIDTH
camera_properties.height = HEIGHT
camera_properties.enable_tensors = True

# rigid shape material properties
shape_props = [gymapi.RigidShapeProperties()]
shape_props[0].friction = -1
shape_props[0].rolling_friction = 1
shape_props[0].torsion_friction = 1
shape_props[0].compliance = 0
shape_props[0].restitution = 0
shape_props1 = [gymapi.RigidShapeProperties()]
shape_props1[0].friction = 0.2
shape_props1[0].rolling_friction = -1
shape_props1[0].torsion_friction = -1
shape_props1[0].compliance = 0
shape_props1[0].restitution = 0
shape_props2 = [gymapi.RigidShapeProperties()]
shape_props2[0].friction = 1
shape_props2[0].rolling_friction = -1
shape_props2[0].torsion_friction = -1
shape_props2[0].compliance = 0
shape_props2[0].restitution = 0

'''
generate WFC maps
'''
PCGWorker_ = PCGWorker(9,9)
seed = PCGWorker_.generate()
seeds = [seed]
for i in range(0,num_matrix_envs-1):
    seed_ = PCGWorker_.mutate(seeds[-1],81)
    seeds.append(seed_)
building_block_pointer = [0,0,0,0,0,0,0,0,0,0,0,0,0]    # 0-4,brick; 5-8,corners; 9-12 ,ramps;

gym = gymapi.acquire_gym()

sim = gym.create_sim(args.compute_device_id, args.graphics_device_id, gymapi.SIM_PHYSX, sim_params)

print("args.graphics_device_id : ",args.graphics_device_id)
print("args.compute_device_id : ",args.compute_device_id)


# create the ground plane
gym.add_ground(sim, plane_params)

capsule_asset = gym.create_capsule(sim, 1, 1, asset_options)

asset_gray_cube = gym.load_asset(sim, asset_root, pt_gray_cube,asset_options1)
asset_blue_cube = gym.load_asset(sim, asset_root, pt_blue_cube,asset_options1)
asset_red_cube = gym.load_asset(sim, asset_root, pt_red_cube,asset_options1)
asset_white_cube = gym.load_asset(sim, asset_root, pt_white_cube,asset_options1)
asset_yellow_cube = gym.load_asset(sim, asset_root, pt_yellow_cube,asset_options1)
asset_blue_ramp = gym.load_asset(sim, asset_root, pt_blue_ramp,asset_options1)
asset_red_ramp = gym.load_asset(sim, asset_root, pt_red_ramp,asset_options1)
asset_white_ramp = gym.load_asset(sim, asset_root, pt_white_ramp,asset_options1)
asset_yellow_ramp = gym.load_asset(sim, asset_root, pt_yellow_ramp,asset_options1)

box_asset = gym.create_box(sim, 1, 1, 1, asset_options)

if not headless:
    cam_props = gymapi.CameraProperties()
    viewer = gym.create_viewer(sim, cam_props)

envs = []
camera_handles = []
actor_handles = []
attractor_handles = []

for n in range(num_matrix_envs):
    # create environments
    env_ = gym.create_env(sim, env_lower, env_upper, 4)
    envs.append(env_)

    n_actors = 0
    # create actors based on WFC seed
    for i in range(0,9):
        for j in range(0,9):
            tile_ = seeds[n].wave_oriented[j*9+i][0][0]
            rot = seeds[n].wave_oriented[j*9+i][0][1]
            pose_ij = gymapi.Transform()
            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 0*side_l)
            pose_ij.r = gymapi.Quat(0,0,0,1)
            actor = gym.create_actor(env_, asset_gray_cube, pose_ij, "gray_cube" + str(i) + str(j), n, 0)
            # cubes
            if tile_<=5:
                if tile_==1:
                    pass
                elif tile_==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                elif tile_==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                elif tile_==4:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                elif tile_==5:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    actor = gym.create_actor(env_, asset_white_cube, pose_ij, "white_cube" + str(i) + str(j), n, 0)
            # corners
            elif tile_==6:
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            elif tile_==7:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            elif tile_==8:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            elif tile_==9:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            # ramps
            elif tile_==10:
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
            elif tile_==11:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
            elif tile_==12:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
            elif tile_==13:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)

    # create capsule_asset actor in the environment
    pose = gymapi.Transform()
    pose.p = gymapi.Vec3(15, 15, 30)
    pose.r = gymapi.Quat(0,1,0,1)
    cap_handle = gym.create_actor(env_, capsule_asset, pose, 'capsule_asset', n, 0)
    gym.set_actor_rigid_shape_properties(env_, cap_handle, shape_props2)

    gym.set_actor_scale(env_, cap_handle, 0.4)
    actor_handles.append(cap_handle)

    attractor_properties_ = gymapi.AttractorProperties()
    attractor_properties_.stiffness = 5e10
    attractor_properties_.damping = 0
    attractor_properties_.axes = gymapi.AXIS_ROTATION
    attractor_properties_.target = pose
    attractor_properties_.rigid_handle = cap_handle
    attractor_handle_ = gym.create_rigid_body_attractor(env_, attractor_properties_)
    gym.set_rigid_body_color(env_, cap_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color_red)
    
    h1 = gym.create_camera_sensor(env_, camera_properties)
    camera_offset = gymapi.Vec3(-2, 0, 0)
    camera_rotation = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 0, 1), 0.4*math.pi)
    body_handle = gym.get_actor_rigid_body_handle(env_, cap_handle, 0)
    gym.attach_camera_to_body(h1, env_, body_handle, gymapi.Transform(camera_offset, camera_rotation), gymapi.FOLLOW_TRANSFORM)
    camera_handles.append(h1)
    attractor_handles.append(attractor_handle_)

    actor_count = gym.get_actor_count(env_)
    # create blank actors
    num_blk_act = max_actors_num_per_env - actor_count
    for i in range(num_blk_act):
        pose = gymapi.Transform()
        pose.p = gymapi.Vec3(0, 0, 0)
        pose.r = gymapi.Quat(0,1,0,1)
        actor = gym.create_actor(env_, box_asset, pose, 'box_asset' + str(i), n, 0)
        gym.set_actor_scale(env_, actor, 0.000001)


# a must-have operation 
gym.prepare_sim(sim)

frame_count = 0

if not headless:
    cam_pos = gymapi.Vec3(0.0, 15.0, 15.0)
    cam_target = gymapi.Vec3(1.5, 4, 0)
    gym.viewer_camera_look_at(viewer, None, cam_pos, cam_target)

# time test
start_time = time.time()
last_frame_cnt = 0

# Create helper geometry used for visualization
# Create an wireframe axis
axes_geom = gymutil.AxesGeometry(2)

facing = 0
facing_step = 0.25

# local coordinate system
verts = np.empty((3, 1), gymapi.Vec3.dtype)
verts[0][0] = (1, 0, 0)
verts[1][0] = (0, 1, 0)
verts[2][0] = (0, 0, 1)

# store initial status for all envs
sim_status = []
for i in range(num_matrix_envs):
    env_status = np.copy(gym.get_env_rigid_body_states(envs[i], gymapi.STATE_ALL))
    sim_status.append(env_status)

while True:

    if not headless:
        if gym.query_viewer_has_closed(viewer):
            break
        # clear lines from previous frame
        gym.clear_lines(viewer)

    '''
    step begin
    '''
    # step the physics
    gym.simulate(sim)
    gym.fetch_results(sim, True)

    # if not headless:
    gym.step_graphics(sim)

    # render the camera sensors
    gym.render_all_camera_sensors(sim)

    for i in range(num_matrix_envs):
        rgb_image = gym.get_camera_image(sim, envs[i], camera_handles[i], gymapi.IMAGE_COLOR)

        if not headless :
            # create an opencv image of size (HEIGHT,WIDTH , 4)
            img = np.zeros((HEIGHT,WIDTH, 4), np.uint8)
            img[:, :, :4] = rgb_image.reshape((HEIGHT,WIDTH, 4))
            # rotate 90 degrees counter-clockwise
            img = np.rot90(img, -1)
            # transform the image from RGBA to RGB
            img = img[:, :, :3]
            # transform the image from RGB to BGR
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

            '''
            render
            '''
            cv2.imshow('rgb', img)
            action = cv2.waitKey(1)

            '''
            step : action
            '''
            if action >0 :
                print(action)
                body_states = gym.get_actor_rigid_body_states(envs[i], actor_handles[i], gymapi.STATE_ALL)
                # transform global coordinate system to local coordinate system
                body_ = gym.get_actor_rigid_body_handle(envs[i], actor_handles[i], 0)
                body_t = gym.get_rigid_transform(envs[i], body_)
                body_t.p = gymapi.Vec3(0,0,0)
                verts_ = body_t.transform_points(verts)
                # print(verts_)

                if action == 32:
                    body_states['vel']['linear'].fill((0,0,spd))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 119:
                    body_states['vel']['linear'].fill((spd * verts_[1][0][0],spd * verts_[1][0][1],spd * verts_[1][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 115:
                    body_states['vel']['linear'].fill((-spd * verts_[1][0][0],-spd * verts_[1][0][1],-spd * verts_[1][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 97:
                    body_states['vel']['linear'].fill((-spd * verts_[2][0][0],-spd * verts_[2][0][1],-spd * verts_[2][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 100:
                    body_states['vel']['linear'].fill((spd * verts_[2][0][0],spd * verts_[2][0][1],spd * verts_[2][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 107:
                    facing -= facing_step
                    attractor_pose = gymapi.Transform()
                    attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
                    attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing*math.pi)
                    gym.set_attractor_target(envs[i], attractor_handles[i], attractor_pose)
                elif action == 108:
                    facing += facing_step
                    attractor_pose = gymapi.Transform()
                    attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
                    attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing*math.pi)
                    gym.set_attractor_target(envs[i], attractor_handles[i], attractor_pose)
                # break !!
                elif action == 111:
                    body_states['vel']['linear'].fill((0,0,0))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)

                # reset !!
                elif action == 114:
                    print("reset")
                    gym.set_env_rigid_body_states(envs[i], sim_status[i], gymapi.STATE_ALL)

    '''
    step end
    '''

    '''
    render
    '''
    # draw camera entities
    if not headless:
        for i in range(num_matrix_envs):
            body_states = gym.get_camera_transform(sim, envs[i], camera_handles[i])
            draw_camera(gym, viewer,envs[i],body_states)
    if not headless:
        # update the viewer
        gym.draw_viewer(viewer, sim, True)

    frame_count += 1
    end_time = time.time()
    if (end_time - start_time) > 1:
        print("FPS: %.2f" % ((frame_count - last_frame_cnt)*num_matrix_envs))
        last_frame_cnt = frame_count
        start_time = time.time()

if not headless:
    gym.destroy_viewer(viewer)
    
gym.destroy_sim(sim)
